#include "kheap.h"
#include "mm.h"
#include "stdio.h"
#include "kernel.h"

static block_head* heap_start;
static block_head* heap_end;
static u32 heap_total_size;

// expands the heap 4kb at a time
static void expand_heap();

u32 heap_size() {
  return heap_total_size;
}

void print_heap() {
  block_head* cur = heap_start;

  while(cur != END_OF_HEAP) {
    if(cur->magic != HEAP_MAGIC) {
      kprintf("corrupted block at %p\n", (u32)cur);
    }

    else {
      kprintf("block at: %p size: %d", (u32)cur + sizeof(block_head), cur->size);
      
      if(cur->free)
	kprintf(" free\n");
      else
	kprintf(" used\n");
    }
    cur = cur->next;
  }
  kprintf("\n");
  return;
}

void init_heap(u32 start, u32 sz) {
  heap_start = (block_head*)start;
  heap_start->magic = HEAP_MAGIC;
  heap_start->size = sz - sizeof(block_head) - sizeof(block_head);
  heap_start->free = TRUE;
  heap_start->prev = 0;
  heap_start->next = (block_head*)(start + sz - sizeof(block_head));
  heap_start->next->size = 0;
  heap_start->next->free = FALSE;
  heap_start->next->magic = HEAP_MAGIC;
  heap_start->next->next = END_OF_HEAP;
  heap_start->next->prev = heap_start;
  heap_end = heap_start->next;
  heap_total_size = sz;
}

u32* kmalloc(u32 sz) {
  block_head* cur = heap_start;
  block_head* next;
  block_head* prev = heap_start->prev;

  // go through all the blocks to see if we can find a free one
  // that is the right size. if next is 0, we are at end of heap
  while(cur->next != END_OF_HEAP) {
    
    // if we find one, expand the next block backwards 
    // to take up the extra slack
    if(cur->size >= (sz + sizeof(block_head)) && cur->free == TRUE) {
      next = cur->next;
      cur->prev = prev;
      cur->next = (block_head*)((u32)cur + sizeof(block_head) + sz);
      cur->next->size = next->size + cur->size - sz;
      cur->next->free = TRUE;
      cur->next->next = next;
      cur->next->prev = cur;
      cur->next->magic = HEAP_MAGIC;
      cur->free = FALSE;
      cur->size = sz;
      return (u32*)((u32)cur + sizeof(block_head));
    }

    prev = cur;
    cur = cur->next;
  }

  // expand heap and try again
  expand_heap();
  return kmalloc(sz);
}

u32* kmalloc_a(u32 sz) {
  block_head* cur = heap_start;
  block_head* next;
  block_head* prev = heap_start->prev;
  u32 tmp, orig_size;  

  while(cur->next != END_OF_HEAP) {
    if(cur->size >= (sz + sizeof(block_head) + (0x1000 - ((u32)cur & 0x00000fff)))
		     && cur->free == TRUE) {
      tmp = (0x1000 + (u32)cur) & 0xFFFFF000;
      next = cur->next;
      orig_size = cur->size;
      cur->size = (tmp - (u32)cur) - sizeof(block_head);
      cur->next = (block_head*)(tmp - sizeof(block_head));
      cur->next->size = sz;
      cur->next->free = FALSE;
      cur->next->magic = HEAP_MAGIC;
      cur->next->prev = cur;
      cur->next->next = (block_head*)(tmp + sz);
      cur->next->next->magic = HEAP_MAGIC;
      cur->next->next->size = orig_size - sz - cur->size;
      cur->next->next->free = TRUE;
      cur->next->next->prev = cur->next;
      cur->next->next->next = next;
      return (u32*)((u32)cur->next + sizeof(block_head));
    }

    prev = cur;
    cur = cur->next;
  }

  // expand heap and try again
  expand_heap();
  return kmalloc_a(sz);
}

// frees a block, and combines it with prev and next blocks if they are free
void kfree(u32* block) {
  block_head* cur = (block_head*)((u32)block - sizeof(block_head));
  cur->free = TRUE;

  if(cur->next->free == TRUE) {
    cur->size = cur->size + cur->next->size;// + sizeof(block_head);
    cur->next = cur->next->next;
    cur->next->prev = cur;
  }

  if(cur->prev->free == TRUE && cur->prev != 0) {
    cur->prev->size = cur->prev->size + cur->size;// + sizeof(block_head);
    cur->next->prev = cur->prev;
    cur->prev->next = cur->next;
  }
  
}

// expand the heap by one page and fix up the heap data struct
static void expand_heap() {
  pte_t* table;
  u32 offset = PTE_OFFSET((u32)heap_end) + 1;

  // if current table is full, map another table for the heap
  if(offset > 0x3ff) {
    kprintf("We have a problem! expand_heap()");
    offset = 0;
  }

  table = (pte_t*)PG_FRAME((u32)kernel_dir->tables[PDE_OFFSET(HEAP_START)]);
  table->pages[offset] = get_frame() | PG_PRES | PG_SUPER | PG_WRITE;
  // alter current end to accomadate new space
  heap_end->size = PG_SIZE - sizeof(block_head);
  heap_end->next = (block_head*)((u32)heap_end + PG_SIZE);
  heap_end->free = TRUE;
  heap_end->next->prev = heap_end;
  // now move the end forward one page
  heap_end = heap_end->next;
  heap_end->magic = HEAP_MAGIC;
  heap_end->size = 0;
  heap_end->next = END_OF_HEAP;
  heap_end->free = FALSE;
  heap_total_size += PG_SIZE;

  // last thing to do is check to see if this is second time in a row we expanded
  // the heap. if so, coalesce the blocks into a big block. this will only occur
  // if someone ever requests something bigger than 4kb, which should never happen
  // in the kernel, but better safe than sorry
  if(heap_end->prev->prev->free == TRUE) {
    heap_end->prev = heap_end->prev->prev;
    heap_end->prev->next = heap_end;
    heap_end->prev->size += PG_SIZE;
  } 

}
